#include <pnm.h>
#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <limits.h>

#ifndef MAX
#define MAX(a, b) ((a)>(b)?(a):(b))
#define MIN(a, b) ((a)<(b)?(a):(b))
#endif

#define BUF_SIZE 33

#define XY_TO_INT(x, y) (((y)<<12)|(x))
#define INT_TO_X(v) ((v)&((1<<12)-1))
#define INT_TO_Y(v) ((v)>>12)
#define INT_PNM_TO_BITMAP(a,b,c) (a | (b>>8) | (c>>16))

int patch_w = 7;
int pm_iters = 5;
int rs_max = INT_MAX;

/* Convert an int a to a binary representation into a char buffer */

char *int2bin(int a, char *buffer, int buf_size) {
    buffer += (buf_size - 1);

    for (int i = 31; i >= 0; i--) {
        *buffer-- = (a & 1) + '0';

        a >>= 1;
    }

    return buffer;
}


/* Measure distance between 2 patches with upper left corners (ax, ay) and (bx, by), 
   terminating early if we exceed a cutoff distance */

int dist(pnm* a, pnm* b, int ax, int ay, int bx, int by, int cutoff) {
  
  int ans = 0;
  int dy;
  
  for(dy=0; dy<patch_w; dy++) {

    int dx;

    int* arowr;
    int* arowg;
    int* arowb;

    int* browr;
    int* browg;
    int* browb;

    arowr = malloc(patch_w*sizeof(int));
    arowg = malloc(patch_w*sizeof(int));
    arowb = malloc(patch_w*sizeof(int));

    browr = malloc(patch_w*sizeof(int));
    browg = malloc(patch_w*sizeof(int));
    browb = malloc(patch_w*sizeof(int));

    int i;

    for(i=0;i<patch_w;i++) {
      arowr[i] = pnm_get_component((*a),ay+dy,ax+i,PnmRed);
      arowg[i] = pnm_get_component((*a),ay+dy,ax+i,PnmGreen);
      arowb[i] = pnm_get_component((*a),ay+dy,ax+i,PnmBlue);

      browr[i] = pnm_get_component((*b),by+dy,bx+i,PnmRed);
      browg[i] = pnm_get_component((*b),by+dy,bx+i,PnmGreen);
      browb[i] = pnm_get_component((*b),by+dy,bx+i,PnmBlue);
    }

    for(dx=0; dx<patch_w; dx++) {

      int dr = arowr[dx] - browr[dx];
      int dg = arowg[dx] - browg[dx];
      int db = arowb[dx] - browb[dx];
      
      ans += dr*dr + dg*dg + db*db;
    }

    free(arowr);
    free(arowg);
    free(arowb);

    free(browr);
    free(browg);
    free(browb);

    if(ans>=cutoff) return cutoff;
  }

  return ans;
}

int dist_default(pnm* a, pnm* b, int ax, int ay, int bx, int by) {
  return dist(a,b,ax,ay,bx,by,INT_MAX);
}

void improve_guess(pnm* a, pnm* b, int ax, int ay, int* xbest, int* ybest, int* dbest,
		   int bx, int by) {

  int d = dist(a,b,ax,ay,bx,by,(*dbest));

  if(d < (*dbest)) {
    (*dbest) = d;
    (*xbest) = bx;
    (*ybest) = by;
  }
}

/* Match image a to image b, returning the nearest neighbor field mapping a => b coords,
   stored in an RGB 24-bit image as (by<<12)|bx */
void patchmatch(pnm* a,pnm* b, char* name_ann, char* name_annd) {

  fprintf(stdout,"==== Creating outputs images ann, annd ====\n\n");

  pnm ann = pnm_new(pnm_get_width(*a),pnm_get_height(*a),PnmRawPpm);
  fprintf(stdout,"  -creating ann: %dx%d\n",pnm_get_width(*a),pnm_get_height(*a));
  pnm annd = pnm_new(pnm_get_width(*a),pnm_get_height(*a),PnmRawPpm);
  fprintf(stdout,"  -creating annd: %dx%d\n",pnm_get_width(*b),pnm_get_height(*b));


  /* Effective width and height (possible upper left corners of patches) */
  int aew = pnm_get_width(*a) - patch_w + 1;
  int aeh = pnm_get_height(*a) - patch_w + 1;

  int bew = pnm_get_width(*b) - patch_w + 1;
  int beh = pnm_get_height(*b) - patch_w + 1;

  int ax, ay;

  srand(time(NULL));

  fprintf(stdout,"Setting values in ann...\n");
  fprintf(stdout,"  -loop x (width vector) from 0 to %d\n",aew);
  fprintf(stdout,"  -loop y (height vector) from 0 to %d\n",aeh);

  for(ay=0; ay<aeh; ay++) {
    for(ax=0; ax<aew; ax++) {      
      int bx = rand()%bew;
      int by = rand()%beh;

      int temp = XY_TO_INT(bx,by);

      pnm_set_component(ann,ay,ax,PnmRed,(temp&255));
      pnm_set_component(ann,ay,ax,PnmGreen,((temp>>8)&255));
      pnm_set_component(ann,ay,ax,PnmBlue,((temp>>16)&255));

      int temp2 = dist_default(a,b,ax,ay,bx,by);

      pnm_set_component(ann,ay,ax,PnmRed,(temp2&255));
      pnm_set_component(ann,ay,ax,PnmGreen,((temp2>>8)&255));
      pnm_set_component(ann,ay,ax,PnmBlue,((temp2>>16)&255));
    }
  }

  int iter;

  fprintf(stdout,"\n==== Patchmatch iteration loop from 0 to %d ====\n\n",pm_iters-1);


  /* In each iteration, improve the NNF, by looping in scanline or reverse-scanline order */
  for(iter=0; iter < pm_iters; iter++) {

    int ystart = 0;
    int yend = aeh;
    int ychange = 1;

    int xstart = 0;
    int xend = aew;
    int xchange = 1;

    if(iter % 2 == 1) {
      
      xstart = xend-1;
      xend = -1;
      xchange = -1;
      
      ystart = yend - 1;
      yend = -1;
      ychange = -1;
    }

    fprintf(stdout,"  -Iteration %d...\n",iter+1);
    /*
    fprintf(stdout,"    -loop x (width vector) from %d to %d\n",xstart,xend);
    fprintf(stdout,"    -loop y (height vector) from %d to %d\n",ystart,yend);
    */

    for(int ay = ystart; ay != yend; ay += ychange) {
      for(int ax = xstart; ax != xend; ax +=xchange) {
	
	int rr = pnm_get_component(ann,ay,ax,PnmRed);
	int gg = pnm_get_component(ann,ay,ax,PnmGreen);
	int bb = pnm_get_component(ann,ay,ax,PnmBlue);

	/* Current (best) guess */
	int v = INT_PNM_TO_BITMAP(rr,gg,bb);

	int xbest = INT_TO_X(v);
	int ybest = INT_TO_Y(v);
	
	rr = pnm_get_component(annd,ay,ax,PnmRed);
	gg = pnm_get_component(annd,ay,ax,PnmGreen);
	bb = pnm_get_component(annd,ay,ax,PnmBlue);

	int dbest = INT_PNM_TO_BITMAP(rr,gg,bb);

	
	/* Propagation: Improve current guess by trying instead correspondences 
	   from left and above (below and right on odd iterations) */
	if ((unsigned) (ax - xchange) < (unsigned) aew) {
	  
	  rr = pnm_get_component(ann,ay,ax - xchange,PnmRed);
	  gg = pnm_get_component(ann,ay,ax - xchange,PnmGreen);
	  bb = pnm_get_component(ann,ay,ax - xchange,PnmBlue);

	  /*
	  char tmpr[BUF_SIZE];
	  tmpr[BUF_SIZE - 1] = '\0';
	  char tmpg[BUF_SIZE];
	  tmpg[BUF_SIZE - 1] = '\0';
	  char tmpb[BUF_SIZE];
	  tmpb[BUF_SIZE - 1] = '\0';
	  
	  int2bin(rr,tmpr,BUF_SIZE - 1);
	  int2bin(gg,tmpg,BUF_SIZE - 1);
	  int2bin(bb,tmpb,BUF_SIZE - 1);

	  fprintf(stdout,"R binary val: %s\n",tmpr);
	  fprintf(stdout,"G binary val: %s\n",tmpg);
	  fprintf(stdout,"B binary val: %s\n",tmpb);

	  */

	  int vp = INT_PNM_TO_BITMAP(rr,gg,bb);

	  /*

	  char tmpv[BUF_SIZE];
	  tmpv[BUF_SIZE - 1] = '\0';
	  
	  int2bin(vp,tmpv,BUF_SIZE - 1);

	  fprintf(stdout,"V binary val: %s\n",tmpv);

	  */

	  int xp = INT_TO_X(vp) + xchange;
	  int yp = INT_TO_Y(vp);

	  if((unsigned) xp < (unsigned) bew)
	    improve_guess(a,b,ax,ay,&xbest,&ybest,&dbest,xp,yp);
	  
	}

	if((unsigned) (ay -ychange) < (unsigned) aeh) {

	  rr = pnm_get_component(ann,ay - ychange,ax,PnmRed);
	  gg = pnm_get_component(ann,ay - ychange,ax,PnmGreen);
	  bb = pnm_get_component(ann,ay - ychange,ax,PnmBlue);

	  int vp = INT_PNM_TO_BITMAP(rr,gg,bb);
	  int xp = INT_TO_X(vp);
	  int yp = INT_TO_Y(vp) + ychange;

	  if((unsigned) yp < (unsigned) beh)
	    improve_guess(a,b,ax,ay,&xbest,&ybest,&dbest,xp,yp);

	  /* Random search: Improve current guess by searching in boxes of 
	     exponentially decreasing size around the current best guess */

	  int rs_start = rs_max;
	  if(rs_start > MAX(pnm_get_width(*b),pnm_get_height(*b))) 
	    rs_start = MAX(pnm_get_width(*b),pnm_get_height(*b));

	  int mag;

	  for(mag = rs_start; mag >= 1; mag /= 2) {

	    /* Sampling window */
	    int xmin = MAX(xbest-mag, 0);
	    int xmax = MIN(xbest+mag+1,bew);
	    int ymin = MAX(ybest-mag,0);
	    int ymax = MIN(ybest+mag+1,beh);
	    int xp = xmin+rand()%(xmax-xmin);
	    int yp = ymin+rand()%(ymax-ymin);
	    improve_guess(a,b,ax,ay,&xbest,&ybest,&dbest,xp,yp);
	  }

	  int tmp = XY_TO_INT(xbest,ybest);
	  	  
	  /*
	    char binarytmp[BUF_SIZE];
	    binarytmp[BUF_SIZE - 1] = '\0';
	    
	    int2bin(tmp,binarytmp,BUF_SIZE - 1);
	    
	    fprintf(stdout,"setting ann %dx%d: %s to 3 int R:%d G:%d B:%d\n",ax,ay,
	    binarytmp,(tmp&255),((tmp>>8)&255),((tmp>>16)&255));
	  */

	  pnm_set_component(ann,ay,ax,PnmRed,tmp&255);
	  pnm_set_component(ann,ay,ax,PnmGreen,(tmp>>8)&255);
	  pnm_set_component(ann,ay,ax,PnmBlue,(tmp>>16)&255);

	  pnm_set_component(annd,ay,ax,PnmRed,dbest&255);
	  pnm_set_component(annd,ay,ax,PnmGreen,(dbest>>8)&255);
	  pnm_set_component(annd,ay,ax,PnmBlue,(dbest>>16)&255);	  
	}
      }
    }
  }
  
  printf("Saving output images\n");
  pnm_save(ann,PnmRawPpm,name_ann);
  pnm_save(annd,PnmRawPpm,name_annd);

  pnm_free(ann);
  pnm_free(annd);
  
}

int main(int argc, char** argv) {

  if(argc != 5) { fprintf(stderr, "pm_minimal a b ann annd\n"
                                   "Given input images a, b outputs nearest neighbor field 'ann' mapping a => b coords, and the squared L2 distance 'annd'\n"
                                   "These are stored as RGB 24-bit images, with a 24-bit int at every pixel. For the NNF we store (by<<12)|bx."); exit(1); }
  printf("Loading input images\n");

  pnm a = pnm_load(argv[1]);
  pnm b = pnm_load(argv[2]);

  printf("Running PatchMatch\n");
  patchmatch(&a,&b,argv[3],argv[4]);  

  pnm_free(a);
  pnm_free(b);
  
  return EXIT_SUCCESS;
}
